from datetime import datetime
import baostock as bs
import pandas as pd
from backtrader.feeds import PandasData


# 创建Backtrader数据源
class BaostockData(PandasData):
    params = (
        ('datetime', 0),
        ('open', 1),
        ('high', 2),
        ('low', 3),
        ('close', 4),
        ('volume', 5),
        ('openinterest', -1),  # 如果没有持仓量数据，设置为-1
    )

    def __init__(self, symbol:str,start_date:datetime, end_date:datetime):
        self.symbol=symbol
        self.start_date = start_date
        self.end_date = end_date
        self.data = self.init_data()
        super(BaostockData, self).__init__(dataname=self.data)
        
    def init_data(self):
        try:
            # 登录Baostock系统
            lg = bs.login()
            if not lg.error_code == '0':
                raise Exception('Login failed with error code: ' + lg.error_code)

            # 获取股票历史K线数据
            rs = bs.query_history_k_data_plus(self.symbol,
                                              "date,code,open,high,low,close,volume,amount",
                                              start_date=self.start_date.strftime('%Y-%m-%d'),
                                              end_date=self.end_date.strftime('%Y-%m-%d'),
                                              frequency="d", adjustflag="3")

            # 打印结果集
            data_list = []
            while (rs.error_code == '0') & rs.next():
                data_list.append(rs.get_row_data())
            data = pd.DataFrame(data_list, columns=rs.fields)

            # 登出系统
            bs.logout()

            # 调整列顺序以符合Backtrader的要求
            data = data[['date', 'open', 'high', 'low', 'close', 'volume']]
            data['date'] = pd.to_datetime(data['date'])

            # 数据验证
            if data.empty:
                raise ValueError('No data retrieved from the query.')

            return data

        except Exception as e:
            print(f'Error occurred: {e}')
            return pd.DataFrame()

# 使用示例
if __name__ == '__main__':
    start_date = datetime(2020, 1, 1)
    end_date = datetime(2020, 12, 31)
    data_feed = BaostockData(start_date=start_date, end_date=end_date)